from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QTableWidget, 
    QTableWidgetItem, QPushButton, QLabel, QHeaderView, 
    QMessageBox, QDialog, QFormLayout, QLineEdit, QTextEdit,
    QDialogButtonBox, QGroupBox, QSplitter, QComboBox
)
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QAction, QCursor

import json
from typing import Dict, List, Any, Optional

from core.chroma_client import ChromaClient


class MetadataDialog(QDialog):
    """元数据键值对编辑对话框"""
    
    def __init__(self, parent=None, key="", value=""):
        """
        初始化元数据对话框
        
        Args:
            parent: 父窗口
            key: 键名
            value: 值
        """
        super().__init__(parent)
        self.setWindowTitle("编辑元数据")
        self.resize(400, 150)
        
        layout = QVBoxLayout()
        
        form_layout = QFormLayout()
        
        self.key_edit = QLineEdit(key)
        form_layout.addRow("键名：", self.key_edit)
        
        self.value_edit = QLineEdit(str(value))
        form_layout.addRow("值:", self.value_edit)
        
        layout.addLayout(form_layout)
        
        # 按钮
        button_box = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | 
            QDialogButtonBox.StandardButton.Cancel
        )
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addWidget(button_box)
        
        self.setLayout(layout)
    
    def get_key(self):
        """获取键名"""
        return self.key_edit.text().strip()
    
    def get_value(self):
        """获取值"""
        value = self.value_edit.text().strip()
        # 尝试转换为数字或布尔值
        if value.lower() == "true":
            return True
        elif value.lower() == "false":
            return False
        try:
            if "." in value:
                return float(value)
            else:
                return int(value)
        except ValueError:
            return value


class DocumentDialog(QDialog):
    """文档编辑对话框"""
    
    def __init__(self, parent=None, document=None, metadata=None, doc_id=None):
        """
        初始化文档对话框
        
        Args:
            parent: 父窗口
            document: 文档内容
            metadata: 文档元数据
            doc_id: 文档ID
        """
        super().__init__(parent)
        self.document = document
        self.metadata = metadata or {}
        self.doc_id = doc_id
        
        self.setWindowTitle("编辑文档" if doc_id else "添加文档")
        self.resize(700, 500)
        
        self._init_ui()
        
        # 如果提供了文档，则填充表单
        if document is not None:
            self.document_edit.setPlainText(document)
            
        if metadata:
            self._fill_metadata_table(metadata)
    
    def _init_ui(self):
        """初始化UI组件"""
        layout = QVBoxLayout()
        
        # 文档内容
        doc_group = QGroupBox("文档内容")
        doc_layout = QVBoxLayout()
        
        self.document_edit = QTextEdit()
        doc_layout.addWidget(self.document_edit)
        doc_group.setLayout(doc_layout)
        
        # 元数据
        metadata_group = QGroupBox("元数据")
        metadata_layout = QVBoxLayout()
        
        # 元数据表格
        self.metadata_table = QTableWidget()
        self.metadata_table.setColumnCount(3)
        self.metadata_table.setHorizontalHeaderLabels(["键", "值", "操作"])
        self.metadata_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
        self.metadata_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        self.metadata_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
        
        metadata_layout.addWidget(self.metadata_table)
        
        # 元数据操作按钮
        btn_layout = QHBoxLayout()
        
        self.add_metadata_btn = QPushButton("添加键值对")
        self.add_metadata_btn.clicked.connect(self._on_add_metadata)
        btn_layout.addWidget(self.add_metadata_btn)
        
        btn_layout.addStretch()
        
        metadata_layout.addLayout(btn_layout)
        metadata_group.setLayout(metadata_layout)
        
        # 按钮
        button_box = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Ok | 
            QDialogButtonBox.StandardButton.Cancel
        )
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        
        # 添加到主布局
        layout.addWidget(doc_group, 3)  # 文档内容占3份
        layout.addWidget(metadata_group, 2)  # 元数据占2份
        layout.addWidget(button_box)
        
        self.setLayout(layout)
    
    def _fill_metadata_table(self, metadata):
        """
        填充元数据表格
        
        Args:
            metadata: 元数据字典
        """
        self.metadata_table.setRowCount(len(metadata))
        
        for i, (key, value) in enumerate(metadata.items()):
            # 键
            key_item = QTableWidgetItem(key)
            self.metadata_table.setItem(i, 0, key_item)
            
            # 值
            value_item = QTableWidgetItem(str(value))
            self.metadata_table.setItem(i, 1, value_item)
            
            # 操作按钮
            btn_widget = QWidget()
            btn_layout = QHBoxLayout()
            btn_layout.setContentsMargins(4, 4, 4, 4)
            btn_layout.setSpacing(8)  # 增加按钮之间的间距
            
            edit_btn = QPushButton("编辑")
            edit_btn.setFixedWidth(60)  # 固定宽度
            edit_btn.setProperty("row", i)
            edit_btn.clicked.connect(self._on_edit_metadata)
            
            delete_btn = QPushButton("删除")
            delete_btn.setFixedWidth(60)  # 固定宽度
            delete_btn.setProperty("row", i)
            delete_btn.clicked.connect(self._on_delete_metadata)
            
            btn_layout.addWidget(edit_btn)
            btn_layout.addWidget(delete_btn)
            btn_widget.setLayout(btn_layout)
            
            self.metadata_table.setCellWidget(i, 2, btn_widget)
    
    def _on_add_metadata(self):
        """添加元数据按钮点击处理"""
        dialog = MetadataDialog(self)
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            key = dialog.get_key()
            value = dialog.get_value()
            
            if not key:
                QMessageBox.warning(self, "验证错误", "键名不能为空")
                return
            
            # 检查是否已存在相同的键
            for i in range(self.metadata_table.rowCount()):
                if self.metadata_table.item(i, 0).text() == key:
                    QMessageBox.warning(self, "验证错误", f"键名 '{key}' 已存在")
                    return
            
            # 添加到表格
            row = self.metadata_table.rowCount()
            self.metadata_table.setRowCount(row + 1)
            
            # 键
            key_item = QTableWidgetItem(key)
            self.metadata_table.setItem(row, 0, key_item)
            
            # 值
            value_item = QTableWidgetItem(str(value))
            self.metadata_table.setItem(row, 1, value_item)
            
            # 操作按钮
            btn_widget = QWidget()
            btn_layout = QHBoxLayout()
            btn_layout.setContentsMargins(4, 4, 4, 4)
            btn_layout.setSpacing(8)  # 增加按钮之间的间距
            
            edit_btn = QPushButton("编辑")
            edit_btn.setFixedWidth(60)  # 固定宽度
            edit_btn.setProperty("row", row)
            edit_btn.clicked.connect(self._on_edit_metadata)
            
            delete_btn = QPushButton("删除")
            delete_btn.setFixedWidth(60)  # 固定宽度
            delete_btn.setProperty("row", row)
            delete_btn.clicked.connect(self._on_delete_metadata)
            
            btn_layout.addWidget(edit_btn)
            btn_layout.addWidget(delete_btn)
            btn_widget.setLayout(btn_layout)
            
            self.metadata_table.setCellWidget(row, 2, btn_widget)
    
    def _on_edit_metadata(self):
        """编辑元数据按钮点击处理"""
        sender = self.sender()
        row = sender.property("row")
        
        key = self.metadata_table.item(row, 0).text()
        value = self.metadata_table.item(row, 1).text()
        
        dialog = MetadataDialog(self, key, value)
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_key = dialog.get_key()
            new_value = dialog.get_value()
            
            if not new_key:
                QMessageBox.warning(self, "验证错误", "键名不能为空")
                return
            
            # 检查是否已存在相同的键（排除当前行）
            for i in range(self.metadata_table.rowCount()):
                if i != row and self.metadata_table.item(i, 0).text() == new_key:
                    QMessageBox.warning(self, "验证错误", f"键名 '{new_key}' 已存在")
                    return
            
            # 更新表格
            key_item = QTableWidgetItem(new_key)
            self.metadata_table.setItem(row, 0, key_item)
            
            value_item = QTableWidgetItem(str(new_value))
            self.metadata_table.setItem(row, 1, value_item)
    
    def _on_delete_metadata(self):
        """删除元数据按钮点击处理"""
        sender = self.sender()
        row = sender.property("row")
        
        # 从表格中删除
        self.metadata_table.removeRow(row)
        
        # 更新后续行的按钮属性
        for i in range(row, self.metadata_table.rowCount()):
            btn_widget = self.metadata_table.cellWidget(i, 2)
            for btn in btn_widget.findChildren(QPushButton):
                btn.setProperty("row", i)
    
    def get_document(self):
        """获取文档内容"""
        return self.document_edit.toPlainText()
    
    def get_metadata(self):
        """获取元数据"""
        metadata = {}
        
        for i in range(self.metadata_table.rowCount()):
            key = self.metadata_table.item(i, 0).text()
            value_text = self.metadata_table.item(i, 1).text()
            
            # 尝试转换为适当的类型
            if value_text.lower() == "true":
                value = True
            elif value_text.lower() == "false":
                value = False
            else:
                try:
                    if "." in value_text:
                        value = float(value_text)
                    else:
                        value = int(value_text)
                except ValueError:
                    value = value_text
            
            metadata[key] = value
        
        return metadata


class CollectionTab(QWidget):
    """集合标签页"""
    
    def __init__(self, client: ChromaClient, collection_name: str, parent=None):
        """
        初始化集合标签页
        
        Args:
            client: ChromaDB客户端
            collection_name: 集合名称
            parent: 父窗口
        """
        super().__init__(parent)
        self.client = client
        self.collection_name = collection_name
        self.collection = client.get_collection(collection_name)
        
        self._init_ui()
        self.load_documents()
    
    def _init_ui(self):
        """初始化UI组件"""
        main_layout = QVBoxLayout()
        
        # 顶部信息
        info_layout = QHBoxLayout()
        
        self.info_label = QLabel(f"集合： {self.collection_name}")
        info_layout.addWidget(self.info_label)
        
        info_layout.addStretch()
        
        # 操作按钮
        button_layout = QHBoxLayout()
        
        self.add_btn = QPushButton("添加文档")
        self.add_btn.clicked.connect(self._on_add_document)
        button_layout.addWidget(self.add_btn)
        
        self.refresh_btn = QPushButton("刷新")
        self.refresh_btn.clicked.connect(self.load_documents)
        button_layout.addWidget(self.refresh_btn)
        
        info_layout.addLayout(button_layout)
        
        # 文档表格
        self.documents_table = QTableWidget()
        self.documents_table.setColumnCount(4)
        self.documents_table.setHorizontalHeaderLabels(["ID", "文档内容", "元数据", "操作"])
        self.documents_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
        self.documents_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        self.documents_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
        self.documents_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
        self.documents_table.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.documents_table.customContextMenuRequested.connect(self._on_table_context_menu)
        
        # 查询面板
        query_group = QGroupBox("查询")
        query_layout = QVBoxLayout()
        
        query_form_layout = QHBoxLayout()
        
        self.query_edit = QLineEdit()
        self.query_edit.setPlaceholderText("输入查询文本...")
        query_form_layout.addWidget(self.query_edit)
        
        self.n_results_combo = QComboBox()
        self.n_results_combo.addItems(["1", "3", "5", "10", "20"])
        self.n_results_combo.setCurrentText("5")
        query_form_layout.addWidget(QLabel("结果数:"))
        query_form_layout.addWidget(self.n_results_combo)
        
        self.query_btn = QPushButton("查询")
        self.query_btn.clicked.connect(self._on_query)
        query_form_layout.addWidget(self.query_btn)
        
        query_layout.addLayout(query_form_layout)
        
        # 查询结果表格
        self.results_table = QTableWidget()
        self.results_table.setColumnCount(4)
        self.results_table.setHorizontalHeaderLabels(["ID", "文档内容", "元数据", "相似度"])
        self.results_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
        self.results_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        self.results_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
        self.results_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
        
        query_layout.addWidget(self.results_table)
        query_group.setLayout(query_layout)
        
        # 创建分割器
        splitter = QSplitter(Qt.Orientation.Vertical)
        
        # 文档面板
        docs_widget = QWidget()
        docs_layout = QVBoxLayout()
        docs_layout.addLayout(info_layout)
        docs_layout.addWidget(self.documents_table)
        docs_widget.setLayout(docs_layout)
        
        splitter.addWidget(docs_widget)
        splitter.addWidget(query_group)
        splitter.setSizes([500, 300])
        
        # 添加到主布局
        main_layout.addWidget(splitter)
        
        self.setLayout(main_layout)
    
    def load_documents(self):
        """加载集合中的文档"""
        try:
            # 获取集合中的所有文档
            result = self.client.get(self.collection_name)
            
            # 更新表格
            self.documents_table.setRowCount(len(result['ids']))
            
            for i, (doc_id, document, metadata) in enumerate(zip(result['ids'], result['documents'], result['metadatas'])):
                # ID
                id_item = QTableWidgetItem(doc_id)
                self.documents_table.setItem(i, 0, id_item)
                
                # 文档内容
                doc_item = QTableWidgetItem(str(document)[:100] + ("..." if len(str(document)) > 100 else ""))
                doc_item.setToolTip(str(document))
                self.documents_table.setItem(i, 1, doc_item)
                
                # 元数据
                try:
                    metadata_str = json.dumps(metadata, ensure_ascii=False)[:100]
                    if len(json.dumps(metadata, ensure_ascii=False)) > 100:
                        metadata_str += "..."
                except:
                    metadata_str = str(metadata)[:100]
                    if len(str(metadata)) > 100:
                        metadata_str += "..."
                
                metadata_item = QTableWidgetItem(metadata_str)
                metadata_item.setToolTip(str(metadata))
                self.documents_table.setItem(i, 2, metadata_item)
                
                # 操作按钮
                btn_widget = QWidget()
                btn_layout = QHBoxLayout()
                btn_layout.setContentsMargins(4, 4, 4, 4)
                btn_layout.setSpacing(8)  # 增加按钮之间的间距
                
                edit_btn = QPushButton("编辑")
                edit_btn.setFixedWidth(60)  # 固定宽度
                edit_btn.setProperty("row", i)
                edit_btn.setProperty("id", doc_id)
                edit_btn.clicked.connect(self._on_edit_document)
                
                delete_btn = QPushButton("删除")
                delete_btn.setFixedWidth(60)  # 固定宽度
                delete_btn.setProperty("row", i)
                delete_btn.setProperty("id", doc_id)
                delete_btn.clicked.connect(self._on_delete_document)
                
                btn_layout.addWidget(edit_btn)
                btn_layout.addWidget(delete_btn)
                btn_widget.setLayout(btn_layout)
                
                self.documents_table.setCellWidget(i, 3, btn_widget)
            
            # 更新信息标签
            self.info_label.setText(f"集合: {self.collection_name} (文档数: {len(result['ids'])})")
            
        except Exception as e:
            QMessageBox.warning(self, "加载错误", f"加载文档时出错: {str(e)}")
    
    def _on_add_document(self):
        """添加文档按钮点击处理"""
        dialog = DocumentDialog(self)
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            document = dialog.get_document()
            metadata = dialog.get_metadata()
            
            if metadata is None:
                return
            
            if not document:
                QMessageBox.warning(self, "验证错误", "文档内容不能为空")
                return
            
            try:
                # 添加文档
                self.client.add(
                    collection_name=self.collection_name,
                    documents=[document],
                    metadatas=[metadata]
                )
                
                # 刷新文档列表
                self.load_documents()
                
            except Exception as e:
                QMessageBox.critical(self, "添加错误", f"添加文档时出错: {str(e)}")
    
    def _on_edit_document(self):
        """编辑文档按钮点击处理"""
        sender = self.sender()
        row = sender.property("row")
        doc_id = sender.property("id")
        
        # 获取文档内容和元数据
        document = self.documents_table.item(row, 1).toolTip()
        metadata_str = self.documents_table.item(row, 2).toolTip()
        
        try:
            metadata = json.loads(metadata_str)
        except:
            metadata = {}
        
        dialog = DocumentDialog(self, document, metadata, doc_id)
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_document = dialog.get_document()
            new_metadata = dialog.get_metadata()
            
            if new_metadata is None:
                return
            
            if not new_document:
                QMessageBox.warning(self, "验证错误", "文档内容不能为空")
                return
            
            try:
                # 更新文档
                self.client.update(
                    collection_name=self.collection_name,
                    document=new_document,
                    metadata=new_metadata,
                    ids=doc_id
                )
                
                # 刷新文档列表
                self.load_documents()
                
            except Exception as e:
                QMessageBox.critical(self, "更新错误", f"更新文档时出错: {str(e)}")
    
    def _on_delete_document(self):
        """删除文档按钮点击处理"""
        sender = self.sender()
        row = sender.property("row")
        doc_id = sender.property("id")
        
        # 确认对话框
        reply = QMessageBox.question(
            self, 
            "确认删除", 
            f"确定要删除ID为 '{doc_id}' 的文档吗？",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            try:
                # 删除文档
                self.client.delete(
                    collection_name=self.collection_name,
                    ids=[doc_id]
                )
                
                # 刷新文档列表
                self.load_documents()
                
            except Exception as e:
                QMessageBox.critical(self, "删除错误", f"删除文档时出错: {str(e)}")
    
    def _on_table_context_menu(self, pos):
        """
        表格上下文菜单处理
        
        Args:
            pos: 菜单位置
        """
        item = self.documents_table.itemAt(pos)
        if not item:
            return
        
        row = item.row()
        doc_id = self.documents_table.item(row, 0).text()
        
        menu = QMenu(self)
        
        edit_action = QAction("编辑文档", self)
        edit_action.triggered.connect(lambda: self._on_edit_document_by_row(row, doc_id))
        menu.addAction(edit_action)
        
        delete_action = QAction("删除文档", self)
        delete_action.triggered.connect(lambda: self._on_delete_document_by_row(row, doc_id))
        menu.addAction(delete_action)
        
        menu.exec(QCursor.pos())
    
    def _on_edit_document_by_row(self, row, doc_id):
        """
        通过行编辑文档
        
        Args:
            row: 行索引
            doc_id: 文档ID
        """
        # 获取文档内容和元数据
        document = self.documents_table.item(row, 1).toolTip()
        metadata_str = self.documents_table.item(row, 2).toolTip()
        
        try:
            metadata = json.loads(metadata_str)
        except:
            metadata = {}
        
        dialog = DocumentDialog(self, document, metadata, doc_id)
        
        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_document = dialog.get_document()
            new_metadata = dialog.get_metadata()
            
            if new_metadata is None:
                return
            
            if not new_document:
                QMessageBox.warning(self, "验证错误", "文档内容不能为空")
                return
            
            try:
                # 更新文档
                self.client.update(
                    collection_name=self.collection_name,
                    document=new_document,
                    metadata=new_metadata,
                    ids=doc_id
                )
                
                # 刷新文档列表
                self.load_documents()
                
            except Exception as e:
                QMessageBox.critical(self, "更新错误", f"更新文档时出错： {str(e)}")
    
    def _on_delete_document_by_row(self, row, doc_id):
        """
        通过行删除文档
        
        Args:
            row: 行索引
            doc_id: 文档ID
        """
        # 确认对话框
        reply = QMessageBox.question(
            self, 
            "确认删除", 
            f"确定要删除ID为 '{doc_id}' 的文档吗？",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )
        
        if reply == QMessageBox.StandardButton.Yes:
            try:
                # 删除文档
                self.client.delete(
                    collection_name=self.collection_name,
                    ids=[doc_id]
                )
                
                # 刷新文档列表
                self.load_documents()
                
            except Exception as e:
                QMessageBox.critical(self, "删除错误", f"删除文档时出错: {str(e)}")
    
    def _on_query(self):
        """查询按钮点击处理"""
        query_text = self.query_edit.text().strip()
        
        if not query_text:
            QMessageBox.warning(self, "验证错误", "查询文本不能为空")
            return
        
        try:
            n_results = int(self.n_results_combo.currentText())
            
            # 执行查询
            results = self.client.query(
                collection_name=self.collection_name,
                query_texts=[query_text],
                n_result=n_results
            )
            
            # 更新结果表格
            self.results_table.setRowCount(len(results['ids'][0]))
            
            for i, (doc_id, document, metadata, distance) in enumerate(zip(
                results['ids'][0], 
                results['documents'][0], 
                results['metadatas'][0],
                results['distances'][0] if 'distances' in results else [0] * len(results['ids'][0])
            )):
                # ID
                id_item = QTableWidgetItem(doc_id)
                self.results_table.setItem(i, 0, id_item)
                
                # 文档内容
                doc_item = QTableWidgetItem(str(document)[:100] + ("..." if len(str(document)) > 100 else ""))
                doc_item.setToolTip(str(document))
                self.results_table.setItem(i, 1, doc_item)
                
                # 元数据
                try:
                    metadata_str = json.dumps(metadata, ensure_ascii=False)[:100]
                    if len(json.dumps(metadata, ensure_ascii=False)) > 100:
                        metadata_str += "..."
                except:
                    metadata_str = str(metadata)[:100]
                    if len(str(metadata)) > 100:
                        metadata_str += "..."
                
                metadata_item = QTableWidgetItem(metadata_str)
                metadata_item.setToolTip(str(metadata))
                self.results_table.setItem(i, 2, metadata_item)
                
                # 相似度
                similarity = 1.0 - float(distance) if distance else 1.0
                similarity_item = QTableWidgetItem(f"{similarity:.4f}")
                self.results_table.setItem(i, 3, similarity_item)
            
        except Exception as e:
            QMessageBox.warning(self, "查询错误", f"执行查询时出错: {str(e)}")
